﻿using System;
using System.Collections.Generic;
using System.Diagnostics.Contracts;
using System.Linq;

namespace GBricks.Collections
{
  public static class MultiDictionary
  {
    public static IDictionary<TKey, IList<TValue>> Create<TKey, TValue>(int capacity = 0, IEqualityComparer<TKey> comparer = null) {
      Contract.Ensures(Contract.Result<IDictionary<TKey, IList<TValue>>>() != null);
      return new Dictionary<TKey, IList<TValue>>(capacity, comparer);
    }

    public static IDictionary<TKey, IList<TValue>> ToMultiDictionary<T, TKey, TValue>(this IEnumerable<T> source,
      Func<T, TKey> key, Func<T, TValue> value, IEqualityComparer<TKey> comparer = null) {
      Argument.NotNull(source, "source");
      Contract.Ensures(Contract.Result<IDictionary<TKey, IList<TValue>>>() != null);
      Func<IEnumerable<TValue>, IList<TValue>> create = items => items.ToList();
      return ToMultiDictionary(source, key, value, create, comparer);
    }

    public static IDictionary<TKey, TCollection> ToMultiDictionary<T, TKey, TValue, TCollection>(this IEnumerable<T> source,
      Func<T, TKey> key, Func<T, TValue> value,
      Func<IGrouping<TKey, TValue>, TCollection> create, IEqualityComparer<TKey> comparer = null)
      where TCollection : ICollection<TValue> {
      Argument.NotNull(source, "source");
      Argument.NotNull(create, "create");
      Contract.Ensures(Contract.Result<IDictionary<TKey, TCollection>>() != null);

      var items =
          from item in source
          group value(item) by key(item) into @group
          select new { @group.Key, Items = create(@group), };
      return items.ToDictionary(item => item.Key, item => item.Items, comparer);
    }

    public static IDictionary<TValue, IList<TKey>> Revert<TKey, TValue>(this IDictionary<TKey, TValue> source, IEqualityComparer<TValue> comparer = null) {
      Argument.NotNull(source, "source");
      Contract.Ensures(Contract.Result<IDictionary<TValue, IList<TKey>>>() != null);
      Func<IEnumerable<TKey>, IList<TKey>> create = items => items.ToList();
      return Revert(source, create, comparer);
    }

    public static IDictionary<TValue, TCollection> Revert<TKey, TValue, TCollection>(this IDictionary<TKey, TValue> source,
      Func<IGrouping<TValue, TKey>, TCollection> create, IEqualityComparer<TValue> comparer = null) where TCollection : ICollection<TKey> {
      Argument.NotNull(source, "source");
      Argument.NotNull(create, "create");
      Contract.Ensures(Contract.Result<IDictionary<TValue, TCollection>>() != null);

      var items =
          from item in source
          group item.Key by item.Value into @group
          select new { @group.Key, Items = create(@group), };
      return items.ToDictionary(item => item.Key, item => item.Items, comparer);
    }

    public static TCollection Add<TKey, TValue, TCollection>(this IDictionary<TKey, TCollection> dictionary, TKey key, TValue value, Func<TCollection> create) where TCollection : ICollection<TValue> {
      Argument.NotNull(dictionary, "dictionary");
      Argument.NotNull(create, "create");

      TCollection collection;
      if(!dictionary.TryGetValue(key, out collection)) {
        collection = create();
        dictionary.Add(key, collection);
      } else if(collection == null) {
        collection = create();
        dictionary[key] = collection;
      }//if

      collection.Add(value);
      return collection;
    }

    public static TCollection Add<TKey, TValue, TCollection>(this IDictionary<TKey, TCollection> dictionary, TKey key, TValue value) where TCollection : ICollection<TValue>, new() {
      Argument.NotNull(dictionary, "dictionary");
      return Add(dictionary, key, value, () => new TCollection());
    }

    public static bool Remove<TKey, TValue, TCollection>(this IDictionary<TKey, TCollection> dictionary, TKey key, TValue value) where TCollection : ICollection<TValue> {
      Argument.NotNull(dictionary, "dictionary");

      TCollection collection;
      if(dictionary.TryGetValue(key, out collection) && collection != null && collection.Remove(value)) {
        if(collection.Count == 0) {
          dictionary.Remove(key);
        }//if

        return true;
      }//if

      return false;
    }

    public static bool Contains<TKey, TValue, TCollection>(this IDictionary<TKey, TCollection> dictionary, TKey key, TValue value) where TCollection : ICollection<TValue> {
      Argument.NotNull(dictionary, "dictionary");

      TCollection collection;
      return dictionary.TryGetValue(key, out collection) && collection != null && collection.Contains(value);
    }

    public static IList<TValue> Add<TKey, TValue>(this IDictionary<TKey, IList<TValue>> dictionary, TKey key, TValue value) {
      Argument.NotNull(dictionary, "dictionary");

      IList<TValue> list;
      if(!dictionary.TryGetValue(key, out list)) {
        list = new List<TValue>();
        dictionary.Add(key, list);
      } else if(list == null) {
        list = new List<TValue>();
        dictionary[key] = list;
      }//if

      list.Add(value);
      return list;
    }

    public static IList<TValue> GetValueOrDefault<TKey, TValue>(this IDictionary<TKey, IList<TValue>> dictionary, TKey key) {
      Argument.NotNull(dictionary, "dictionary");

      IList<TValue> list;
      if(!dictionary.TryGetValue(key, out list) || list == null) {
        return new List<TValue>();
      }//if

      return list;
    }

    public static IList<TValue> GetValueOrNull<TKey, TValue>(this IDictionary<TKey, IList<TValue>> dictionary, TKey key) {
      Argument.NotNull(dictionary, "dictionary");

      IList<TValue> list;
      if(!dictionary.TryGetValue(key, out list)) {
        return null;
      }//if

      return list;
    }

    public static bool Remove<TKey, TValue>(this IDictionary<TKey, IList<TValue>> dictionary, TKey key, TValue value) {
      Argument.NotNull(dictionary, "dictionary");

      IList<TValue> list;
      if(dictionary.TryGetValue(key, out list) && list != null && list.Remove(value)) {
        if(list.Count == 0) {
          dictionary.Remove(key);
        }//if

        return true;
      }//if

      return false;
    }

    public static bool Contains<TKey, TValue>(this IDictionary<TKey, IList<TValue>> dictionary, TKey key, TValue value) {
      Argument.NotNull(dictionary, "dictionary");

      IList<TValue> list;
      return dictionary.TryGetValue(key, out list) && list != null && list.Contains(value);
    }

    public static ICollection<TValue> GetValueOrDefault<TKey, TValue>(this IDictionary<TKey, ICollection<TValue>> dictionary, TKey key) {
      Argument.NotNull(dictionary, "dictionary");

      ICollection<TValue> list;
      if(!dictionary.TryGetValue(key, out list) || list == null) {
        return new List<TValue>();
      }//if

      return list;
    }

    public static ICollection<TValue> GetValueOrNull<TKey, TValue>(this IDictionary<TKey, ICollection<TValue>> dictionary, TKey key) {
      Argument.NotNull(dictionary, "dictionary");

      ICollection<TValue> list;
      if(!dictionary.TryGetValue(key, out list)) {
        return null;
      }//if

      return list;
    }
  }
}
